{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Elaborate EKF Example\n",
    "\n",
    "Simple linear Kalman filter on a moving 2D point using factor graphs in GTSAM.\n",
    "This example manually creates all of the needed data structures to show how\n",
    "the Kalman filter works under the hood using factor graphs, but uses a loop\n",
    "to handle the repetitive prediction and update steps.\n",
    "\n",
    "Author: Matt Kielo. Based on the C++ example by Frank Dellaert and Stephen Williams"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "source": [
    "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\n",
    "Atlanta, Georgia 30332-0415\n",
    "All Rights Reserved\n",
    "\n",
    "Authors: Frank Dellaert, et al. (see THANKS for the full author list)\n",
    "\n",
    "See LICENSE for the license information"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/python/gtsam/examples/elaboratePoint2KalmanFilter.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gtsam\n",
    "import numpy as np\n",
    "from gtsam import Point2, noiseModel\n",
    "from gtsam.symbol_shorthand import X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code below basically implements the SRIF (Square-root Information filter version of the EKF) with Cholesky factorization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup containers for linearization points\n",
    "linearization_points = gtsam.Values()\n",
    "\n",
    "# Initialize state x0 at origin\n",
    "x_initial = Point2(0, 0)\n",
    "p_initial = noiseModel.Isotropic.Sigma(2, 0.1)\n",
    "\n",
    "# Add x0 to linearization points\n",
    "linearization_points.insert(X(0), x_initial)\n",
    "\n",
    "# Initial factor graph with prior on X(0)\n",
    "gfg = gtsam.GaussianFactorGraph()\n",
    "ordering = gtsam.Ordering()\n",
    "ordering.push_back(X(0))\n",
    "gfg.add(X(0), p_initial.R(), np.zeros(2), noiseModel.Unit.Create(2))\n",
    "\n",
    "# Common parameters for all steps\n",
    "motion_delta = Point2(1, 0)  # Always move 1 unit to the right\n",
    "process_noise = noiseModel.Isotropic.Sigma(2, 0.1)\n",
    "measurement_noise = noiseModel.Isotropic.Sigma(2, 0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X1 Predict: [1. 0.]\n",
      "X1 Update: [1. 0.]\n",
      "X2 Predict: [2. 0.]\n",
      "X2 Update: [2. 0.]\n",
      "X3 Predict: [3. 0.]\n",
      "X3 Update: [3. 0.]\n",
      "\n",
      "Elaborate Final Covariance (after update):\n",
      " [[0.0193 0.    ]\n",
      " [0.     0.0193]]\n"
     ]
    }
   ],
   "source": [
    "# Current state and conditional\n",
    "current_x = X(0)\n",
    "current_conditional = None\n",
    "current_result = None\n",
    "\n",
    "# Run three predict-update cycles\n",
    "for step in range(1, 4):\n",
    "    # =====================================================================\n",
    "    # Prediction step\n",
    "    # =====================================================================\n",
    "    next_x = X(step)\n",
    "    \n",
    "    # Create new graph with prior from previous step if not the first step\n",
    "    if step > 1:\n",
    "        gfg = gtsam.GaussianFactorGraph()\n",
    "        gfg.add(\n",
    "            current_x,\n",
    "            current_conditional.R(),\n",
    "            current_conditional.d() - current_conditional.R() @ current_result.at(current_x),\n",
    "            current_conditional.get_model()\n",
    "        )\n",
    "    \n",
    "    # Add next state to ordering and create motion model\n",
    "    ordering = gtsam.Ordering()\n",
    "    ordering.push_back(current_x)\n",
    "    ordering.push_back(next_x)\n",
    "    \n",
    "    # Create motion factor and add to graph\n",
    "    motion_factor = gtsam.BetweenFactorPoint2(current_x, next_x, motion_delta, process_noise)\n",
    "    \n",
    "    # Add next state to linearization points if this is the first step\n",
    "    if step == 1:\n",
    "        linearization_points.insert(next_x, x_initial)\n",
    "    else:\n",
    "        linearization_points.insert(next_x, \n",
    "                                    linearization_points.atPoint2(current_x))\n",
    "    \n",
    "    # Add linearized factor to graph\n",
    "    gfg.push_back(motion_factor.linearize(linearization_points))\n",
    "    \n",
    "    # Solve for prediction\n",
    "    prediction_bayes_net = gfg.eliminateSequential(ordering)\n",
    "    next_conditional = prediction_bayes_net.back()\n",
    "    prediction_result = prediction_bayes_net.optimize()\n",
    "    \n",
    "    # Extract and store predicted state\n",
    "    next_predict = linearization_points.atPoint2(next_x) + Point2(prediction_result.at(next_x))\n",
    "    print(f\"X{step} Predict:\", next_predict)\n",
    "    linearization_points.update(next_x, next_predict)\n",
    "    \n",
    "    # =====================================================================\n",
    "    # Update step\n",
    "    # =====================================================================\n",
    "    # Create new graph with prior from prediction\n",
    "    gfg = gtsam.GaussianFactorGraph()\n",
    "    gfg.add(\n",
    "        next_x,\n",
    "        next_conditional.R(),\n",
    "        next_conditional.d() - next_conditional.R() @ prediction_result.at(next_x),\n",
    "        next_conditional.get_model()\n",
    "    )\n",
    "    \n",
    "    # Create ordering for update\n",
    "    ordering = gtsam.Ordering()\n",
    "    ordering.push_back(next_x)\n",
    "    \n",
    "    # Create measurement at correct position\n",
    "    measurement = Point2(float(step), 0.0)\n",
    "    meas_factor = gtsam.PriorFactorPoint2(next_x, measurement, measurement_noise)\n",
    "    \n",
    "    # Add measurement factor to graph\n",
    "    gfg.push_back(meas_factor.linearize(linearization_points))\n",
    "    \n",
    "    # Solve for update\n",
    "    update_bayes_net = gfg.eliminateSequential(ordering)\n",
    "    current_conditional = update_bayes_net.back()\n",
    "    current_result = update_bayes_net.optimize()\n",
    "    \n",
    "    # Extract and store updated state\n",
    "    next_update = linearization_points.atPoint2(next_x) + Point2(current_result.at(next_x))\n",
    "    print(f\"X{step} Update:\", next_update)\n",
    "    linearization_points.update(next_x, next_update)\n",
    "    \n",
    "    # Move to next state\n",
    "    current_x = next_x\n",
    "\n",
    "final_R = current_conditional.R()\n",
    "final_information = final_R.transpose() @ final_R\n",
    "final_covariance = np.linalg.inv(final_information)\n",
    "print(\"\\nElaborate Final Covariance (after update):\\n\", final_covariance)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
